package main

import "fmt"
import "strconv"
import "flag"
import "os"
import "time"
import "bufio" //用于逐行读取文本
import httpclient "github.com/ddliu/go-httpclient"

// URL列表
var urlList chan string = make(chan string,100)
// 检测结果
var result chan string = make(chan string,100)
// 线程数
var thread chan int = make(chan int)
var workThread int

// 关闭程序
var clo chan bool = make(chan bool,100)

// 保存结果
func writeResult() {
    fileName := "result.txt"
	fout,err := os.Create(fileName)
	if err != nil {
		// 文件创建失败
		fmt.Println(fileName + " create failed!")
	}
	defer fout.Close()
	s,ok := <- result
	for ;ok; {
		fout.WriteString(s + "\r\n")
		s,ok = <- result		
	}
	// 等待结果保存完成后通知进程退出
	clo <- true
}

// 根据线程参数启动检测器
func runChecker() {
	t,ok := <- thread
	workThread = t
	if ok {
		for i := 0; i < workThread; i++ {
			go checker(strconv.Itoa(i))
		}
	}
	// 等待线程终止
	for ; <- thread == 0 ; {
		workThread--
		if workThread == 0 {
			// 全部线程已终止，关闭结果写入，退出程序
			close(result)
			break
		}
	}
}

// 检测线程
func checker(threadId string) {
	
	t := time.NewTimer(150 * time.Minute)
	defer t.Stop()
	select {
	case url,ok := <- urlList : 
	for ; ok; {
		fmt.Println("[thread-" + threadId + "] Check:")
		c := httpclient.NewHttpClient(map[int]interface{} {httpclient.OPT_USERAGENT: "my awsome httpclient",})
		response,_ := c.Get(url, nil)
		if response.StatusCode == 200 {
			// URL可访问
			result <- url
		}
		url,ok = <- urlList
	}

	case <- t.C : fmt.Println("检测URL地址超时")
	}

	fmt.Println("[thread-" + threadId + "] end")
	thread <- 0
}

// 读取URL列表文件
func getUrls(fileName string) {
	fin,err := os.Open(fileName)
	defer fin.Close()
	if err != nil {
	        fmt.Println( err)
	        return
	}
	buff := bufio.NewReader(fin)
	for {
		// 逐行读取，没有多余字符
		line,_ , err := buff.ReadLine()
		// 信道超时设置
		t := time.NewTimer(15 * time.Minute)
		defer t.Stop()
		select {
		case urlList <- string(line) : ;
		case <- t.C : fmt.Println("写入URL地址信道超时")
		}
		if err != nil {
			break
		}
	}
	close(urlList)
}

// 处理参数
func processFlag(arg []string) {
	// 线程数
	t,err := strconv.Atoi(arg[1])
	if err != nil {
		// 线程数设置错误
		fmt.Println("'thread' Setting error")
		return
	}
	// 限制最大线程数4096
	if t < 1 {
		t = 1
	}else if t > 4096 {
		t = 4096
	}
	// 传递启动进程数
	thread <- t

	// URLs列表文件
	fileName := arg[0]
	// 传递URL列表
	getUrls(fileName)
	
}

func main() {
	flag.Parse()
	if flag.NArg() != 2 {
		// 参数不合法
		fmt.Println("Usage: checkURL URL.lst threadNumber")
		fmt.Println("Author: Frank Kung <i@kongfanjian.com>")
		return
	}

	// 获取参数
	args := make([]string,0,2)
	for i := 0; i < flag.NArg(); i++ {
		args = append(args,flag.Arg(i))
	}

	// 启动检测线程
	go runChecker()
	// 启动结果写入线程
	go writeResult()
	// 参数处理
	processFlag(args)
	time.Sleep(10000)
	// 等待退出指令
	<- clo
	fmt.Println("Exit")
}